import shutil
import os
import configparser
import paramiko
import socket
import select
import threading
import hashlib


class SSHTunnel:
    def __init__(self, config_file='config.ini', callback=None):

        self.targets = None
        self.server_a_password = None
        self.server_a_username = None
        self.server_a_port = None
        self.server_a_ip = None
        self.config = None
        self.running_local_server = True  # 添加一个标志位
        self.running_handle_channel = True  # 添加一个新的标志位
        self.callback = callback
        self.parse_config(config_file)
        # 创建了一个paramiko.SSHClient对象，用于与远程SSH服务器建立连接
        self.ssh_client = paramiko.SSHClient()
        # 置了SSH客户端的缺失主机密钥策略。paramiko.AutoAddPolicy()策略允许SSH客户端自动接受远程服务器的主机密钥，并将其添加到本地的known_hosts文件中。
        # 这对于首次连接到新服务器时很有用，但在生产环境中可能存在安全风险，因为它可能会使客户端容易受到中间人攻击。在生产环境中，建议使用paramiko.WarningPolicy或自定义策略来确保主机密钥的安全性
        self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.server_threads = []

    @staticmethod
    def calculate_md5(file_path):
        with open(file_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()

    def parse_config(self, config_file):
        msg = f"""Warning in [parse_config]: 
已生成配置文件, 请编辑通目录下的{config_file}后再开始隧道服务
配置文件格式如下: 
[ServerA]: 填写ssh服务器的地址
[Target(序号)]:  target_ip:target_port填写目标地址与端口
[Target(序号)]:  local_listen_ip:local_listen_port 填写本地监听的地址与端口,local_listen_ip默认127.0.0.1"""
        self.config = configparser.ConfigParser()
        if not os.path.exists(config_file):
            # 添加ServerA部分
            self.config['ServerA'] = {
                'ip': '192.168.123.228',
                'port': '22',
                'username': 'root',
                'password': '******'
            }
            # 添加Target1部分
            self.config['Target1'] = {
                'target_ip': '172.168.123.228',
                'target_port': '3306',
                'local_listen_ip': '127.0.0.1',
                'local_listen_port': '3306'
            }
            # 添加Target2部分
            self.config['Target2'] = {
                'target_ip': '172.168.123.228',
                'target_port': '443',
                'local_listen_ip': '127.0.0.1',
                'local_listen_port': '443'
            }
            with open('config.ini', 'w') as configfile:
                self.config.write(configfile)
            raise Exception(msg)
        else:
            if self.calculate_md5(config_file) == "d386fac704b84aeba37b743d62ce8dfe":
                raise Exception(msg)

        self.config.read(config_file)
        self.server_a_ip = self.config.get('ServerA', 'ip', fallback="")
        self.server_a_port = self.config.getint('ServerA', 'port', fallback=0)
        self.server_a_username = self.config.get('ServerA', 'username', fallback="")
        self.server_a_password = self.config.get('ServerA', 'password', fallback="")
        self.targets = []
        for section in self.config.sections():
            if section.startswith('Target'):
                target = {
                    "target_ip": self.config.get(section, 'target_ip', fallback=""),
                    "target_port": self.config.getint(section, 'target_port', fallback=""),
                    "local_listen_ip": self.config.get(section, 'local_listen_ip', fallback=""),
                    "local_listen_port": self.config.getint(section, 'local_listen_port', fallback=""),
                }
                for key, value in target.items():
                    if not value:
                        msg = f"Error in [__init__]: 配置项不能为空: {section}.{key} 为空, 请检查配置文件: {config_file}"
                        raise Exception(msg)
                self.targets.append(target)

    def print_msg(self, msg):
        if self.callback:
            self.callback(msg)
        else:
            print(msg)

    def handle_channel(self, channel, local_conn):
        try:

            while self.running_handle_channel:
                # select.select()函数是一个I/O多路复用函数，它允许程序在多个套接字上等待事件（如可读、可写或异常）
                # select.select()：select.select()函数接受三个参数，分别是要监视的可读套接字列表、可写套接字列表和异常套接字列表。在这个例子中，我们只关心可读套接字，
                # 所以传递了[channel, local_conn]作为第一个参数，而第二个和第三个参数是空列表。
                # 返回值：select.select()函数返回三个列表，分别是可读套接字列表、可写套接字列表和异常套接字列表。在这个例子中，我们只关心可读套接字，所以只处理第一个返回的列表。
                r, _, _ = select.select([channel, local_conn], [], [])
                # if channel in r:：检查SSH通道（channel）是否在可读套接字列表（r）中。如果是，说明通道上有数据可读。
                # data = channel.recv(4096)：从SSH通道上接收最多4096字节的数据。这里使用了recv()方法，它是在套接字上调用的，用于接收数据。
                # if len(data) == 0:：检查接收到的数据长度是否为0。如果是，说明通道已经关闭，我们应该跳出循环。
                # local_conn.send(data)：将从SSH通道接收到的数据发送到本地连接（local_conn）。这里使用了send()方法，它是在套接字上调用的，用于发送数据。
                if channel in r:
                    data = channel.recv(4096)
                    if len(data) == 0:
                        break
                    local_conn.send(data)
                # if local_conn in r:：检查本地连接（local_conn）是否在可读套接字列表（r）中。如果是，说明连接上有数据可读。
                # data = local_conn.recv(4096)：从本地连接上接收最多4096字节的数据。这里使用了recv()方法，它是在套接字上调用的，用于接收数据。
                # if len(data) == 0:：检查接收到的数据长度是否为0。如果是，说明连接已经关闭，我们应该跳出循环。
                # channel.send(data)：将从本地连接接收到的数据发送到SSH通道（channel）。这里使用了send()方法，它是在套接字上调用的，用于发送数据。
                if local_conn in r:
                    data = local_conn.recv(4096)
                    if len(data) == 0:
                        break
                    channel.send(data)
        except Exception as e:
            msg = f"Error in [handle_channel]: {e}"
            self.print_msg(msg)
        finally:
            channel.close()
            local_conn.close()

    def handle_local_connection(self, local_conn, ssh_client, target_ip, target_port):
        try:
            # getsockname()方法返回一个包含套接字本地地址和端口的元组
            # local_listen_ip, local_listen_port = local_conn.getsockname()
            # get_transport() 方法返回一个SSH传输对象，该对象可以用于创建SSH通道、发送和接收数据等操作
            transport = ssh_client.get_transport()
            # (target_ip, target_port) 和 (server_a_ip, 0) 是两个元组，分别表示目标服务器的IP地址和端口号，以及SSH服务器的IP地址和端口号。
            # 在这个例子中，我们使用(server_a_ip, 0)作为SSH服务器的地址和端口，这意味着我们将在本地端口上监听连接，并将它们转发到目标服务器。
            channel = transport.open_channel("direct-tcpip", (target_ip, target_port), (self.server_a_ip, 0))
            thread = threading.Thread(target=self.handle_channel, args=(channel, local_conn))
            thread.start()
        except Exception as e:
            msg = f"Error in [handle_local_connection]: {e}"
            self.print_msg(msg)
            local_conn.close()

    def start_local_server(self, ssh_client, local_listen_ip, local_listen_port, target_ip, target_port):
        try:
            msg = f"正在创建本地监听 {local_listen_ip}:{local_listen_port}"
            self.print_msg(msg)

            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as local_server:
                local_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                local_server.bind((local_listen_ip, local_listen_port))
                local_server.listen(20)
                msg = f"本地监听创建成功 {local_listen_ip}:{local_listen_port}"
                self.print_msg(msg)

                while self.running_local_server:
                    local_conn, _ = local_server.accept()
                    thread = threading.Thread(target=self.handle_local_connection,
                                              args=(local_conn, ssh_client, target_ip, target_port))
                    thread.start()
        except OSError as e:
            msg = f"Error in [start_local_server] 监听失败, 本地端口{local_listen_port}已被占用\n{e}"
            self.print_msg(msg)
            return
        except Exception as e:
            msg = f"Error in [start_local_server]: {e}"
            self.print_msg(msg)
            return

    def start(self):
        # 连接到远程SSH服务器
        try:
            msg = f"正在连接服务器, 请等候, 超时10秒..."
            self.print_msg(msg)
            self.ssh_client.connect(self.server_a_ip, self.server_a_port, self.server_a_username,
                                    self.server_a_password, timeout=10)
            self.print_msg(f"连接服务器{self.server_a_ip}:{self.server_a_port}成功, 正在建立隧道...")
        except Exception as e:
            msg = f"Error in [start] 服务器连接失败{self.server_a_ip}:{self.server_a_port}: {e}"
            self.print_msg(msg)
            return

        for target in self.targets:
            self.print_msg(msg)
            t = threading.Thread(target=self.start_local_server, args=(
                self.ssh_client, target["local_listen_ip"], target["local_listen_port"], target["target_ip"],
                target["target_port"]))
            t.start()
            self.server_threads.append(t)
            msg = f"隧道创建成功 {target['local_listen_ip']}:{target['local_listen_port']} --> {target['target_ip']}:{target['target_port']}"
            self.print_msg(msg)
        # try:
        #     for t in self.server_threads:
        #         t.join()
        # except KeyboardInterrupt:
        #     print("Exiting...")
        # finally:
        #     self.ssh_client.close()

    def stop_all_threads(self):
        self.running_local_server = False
        self.running_handle_channel = False
        for thread in self.server_threads:
            if thread.is_alive():
                thread.join()
        self.ssh_client.close()
        # except Exception as e:
        #     msg = f"未启动隧道转发: {e}"
        #     self.print_msg(msg)


if __name__ == "__main__":
    tunnel = SSHTunnel()
    tunnel.start()
    # 在这里添加代码以便在需要时调用tunnel.stop()来停止SSH隧道
